from data_loader import load_cifar10_data
from model import Cifar10ResNet
from train import train_model
from evaluate import evaluate_model
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR


if __name__ == "__main__":
    train_loader, test_loader = load_cifar10_data()

    model = Cifar10ResNet()

    criterion = nn.CrossEntropyLoss()
    # 使用Adam优化器并添加L2正则化
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.0001)
    # 学习率衰减策略
    scheduler = StepLR(optimizer, step_size=7, gamma=0.1)

    train_model(model, train_loader, criterion, optimizer, scheduler, epochs=20)

    accuracy = evaluate_model(model, test_loader)
    # 保存模型
    torch.save(model.state_dict(), "image_classifier.pth")